# -*- coding: utf-8 -*-
import mindspore
from mindspore import Tensor
import numpy as np
from PIL import Image
import os
import argparse
from mindspore import context
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.train import Model
import importlib

def predict(modulePath, moduleName, ckptName, testImgPath, blackBg, resize):
    # load the network definition
    module = importlib.import_module(modulePath)
    network = getattr(module, moduleName)()
    # load the saved model for evaluation
    param_dict = load_checkpoint(ckptName)
    # load parameter to the network
    load_param_into_net(network, param_dict)
    network.batch_size = 100
    
    model = Model(network)
    # define the input Tensor
    imgList = []
    labelList = []
    '''
    the test images is store like below
    testImgPath
           |-----0
           |     |-----0.png
           |     |-----1.png
           |      -----...
           |-----1
           |     |-----0.png
           |     |-----1.png
           |      -----...
           |-----2
           |     |-----0.png
           |     ...
    '''
    for root, dirs, files in os.walk(testImgPath):
        for file in files:
            # filter
            if 'scan' in file:
                continue
            img=Image.open(os.path.join(root, file)).convert('L').resize((resize, resize), Image.ANTIALIAS)
            img=np.array(img)
            # label is in the dir's name
            labelList.append(int(root[-1]))
            # if the background is white, need to reverse the image to black background 
            if not blackBg:
                img = 255.0 - img
            # rescale images
            img = img / 255.0
            # normalize images
            img = (img * 1 / 0.3081) - 1 * 0.1307 / 0.3081
            # store the image as NCHW
            imgList.append([img])

    x = Tensor(imgList, mindspore.float32)
    y = model.predict(x).asnumpy()
    print('predict:', np.argmax(y, axis=1).tolist())
    print('labels :', labelList)
    print("Accuracy: ", np.sum((np.argmax(y, axis=1) - np.array(labelList))==0)/len(labelList)*100, "%")
    
def str2bool(str):
    return True if str.lower() == 'true' else False    


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='MindSpore Exam Test Tool')
    parser.add_argument('--device_target', type=str, default="CPU", choices=['Ascend', 'GPU', 'CPU'],
                        help='device where the code will be implemented (default: CPU)')
    parser.add_argument('--ckpt_name', type=str, default="mnist.ckpt",
                        help='checkpoint file name, default is mnist.ckpt')
    parser.add_argument('--test_image_path', type=str, default="./test/",
                        help='a dir with subdir of name 0-9 which contain images of 0-9')
    parser.add_argument('--module_path', type=str, default="mnist_tutorial",
                        help='module path that contains the definition of the network')
    parser.add_argument('--module_name', type=str, default="Net",
                        help='network\'s class name')                        
    parser.add_argument('--resize', type=int, default=32,
                        help='resize pixel, height and width')                        
    parser.add_argument('--black_bg', type=str2bool, default=True,
                        help='black background or not')
    args = parser.parse_args()
    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)

    predict(args.module_path, args.module_name, args.ckpt_name, args.test_image_path, args.black_bg, args.resize)
